{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "df0d74f3-487b-4f31-9750-fd9f7d572932",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4cd2e6ad-55c0-431c-971c-9a9208904f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "SECONDS_PER_YEAR = 365 * 24 * 60 * 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1322f0c-bf45-438b-aa36-eafce130147f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apr_to_apy(apr: float, compounds: int) -> float:\n",
    "    if compounds <= 0:\n",
    "        return apr\n",
    "    periodic = apr / compounds\n",
    "    return (1 + periodic) ** compounds - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d86910eb-2a7d-423a-be04-3e6018cc373f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reward_rate : same value as given by `ProtocolStaking:rewardRate()`\n",
    "# num_tokens_per_pool: same values as given by `ProtocolStaking:balanceOf(address(OperatorStaking))` for every eligible OperatorStaking\n",
    "# owner_fees_per_pool: same values as given by `OperatorRewarder:ownerFeeBasisPoints()` for every corresponding OperatorRewarder\n",
    "def compute_native_APR(reward_rate: int, num_tokens_per_pool: list[int], owner_fees_per_pool: list[int]) -> list[float]:\n",
    "    \"\"\"\n",
    "    :param reward_rate: in amount of tokens per second according to the EVM, i.e reward_rate of 1e18 <=> 1 ZAMA token per second\n",
    "    :param num_tokens_per_pool: list of amount of tokens deposited in each OperatorStaking pool\n",
    "    :param owner_fees_per_pool: list of owner fees, one per OperatorStaking pool, in basis points, i.e owner_fee of 10000 <=> 100% fees \n",
    "    \n",
    "    :return: a float corresponding to the PERCENTAGE of _native_ APR i.e not considering direct token donations to OperatorStaking.\n",
    "    It assumes reward_rate, num_tokens_per_pool and list_owner_fees are constants, to get an instantaneous APR value.\n",
    "    \"\"\"\n",
    "    assert len(num_tokens_per_pool)==len(owner_fees_per_pool), \"`num_tokens_per_pool` and `owner_fees_per_pool` length mismatch\"\n",
    "    assert all(0 <= fee <= 10000 for fee in owner_fees_per_pool), \"owner fees must be in [0, 10000]\"\n",
    "    assert all(0 <= num_token for num_token in num_tokens_per_pool), \"number of tokens must be non-negative\"\n",
    "    assert reward_rate >=0, \"`reward_rate` must be non-negative\"\n",
    "\n",
    "    weights = [int(math.sqrt(x)) for x in num_tokens_per_pool]\n",
    "    total_weight = sum(weights)\n",
    "    fee_factors = [1-fee/10000  for fee in owner_fees_per_pool]\n",
    "    reward_rate_per_sec_per_pool = [reward_rate * (pool_weight / total_weight) for pool_weight in weights]\n",
    "    pool_aprs = []\n",
    "    for i in range(len(fee_factors)):\n",
    "        net_reward_per_sec_per_pool = reward_rate_per_sec_per_pool[i] * fee_factors[i]\n",
    "        pool_aprs.append((net_reward_per_sec_per_pool / num_tokens_per_pool[i]) * SECONDS_PER_YEAR * 100)\n",
    "    \n",
    "    return pool_aprs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4802969a-2f8d-4a49-9fa3-8f293a719865",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_native_APY(\n",
    "    reward_rate: int,\n",
    "    num_tokens_per_pool: list[int],\n",
    "    owner_fees_per_pool: list[int],\n",
    "    compounds: int,\n",
    ") -> list[float]:\n",
    "    \"\"\"\n",
    "    Same inputs as compute_native_APR plus `compounds` per year.\n",
    "    Returns APY as PERCENT (e.g. 5.1 for 5.1%).\n",
    "    \"\"\"\n",
    "    aprs_percent = compute_native_APR(reward_rate, num_tokens_per_pool, owner_fees_per_pool)\n",
    "    apys_percent = [\n",
    "        apr_to_apy(apr / 100.0, compounds) * 100.0   # convert % -> fraction -> APY -> %\n",
    "        for apr in aprs_percent\n",
    "    ]\n",
    "    return apys_percent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ea092cbf-f51e-4960-ab2a-a30effd7b7b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "APRs in % for each pool  : [5.991840000000001, 5.991840000000001, 5.991840000000001, 5.991840000000001, 5.991840000000001]\n",
      "APYs in % for same pools : [6.174468300782898, 6.174468300782898, 6.174468300782898, 6.174468300782898, 6.174468300782898]\n"
     ]
    }
   ],
   "source": [
    "print(\"APRs in % for each pool  :\" , compute_native_APR(1e18, 5*[100_000_000*1e18], 5*[500]))\n",
    "print(\"APYs in % for same pools :\", compute_native_APY(1e18, 5*[100_000_000*1e18], 5*[500], 365))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5318b359-36d3-4eec-8b45-d509880699ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "APRs in % for each pool  : [6.307199999999999, 5.991840000000001, 5.991840000000001, 5.991840000000001, 5.991840000000001]\n",
      "APYs in % for same pools : [6.509772041144912, 6.174468300782898, 6.174468300782898, 6.174468300782898, 6.174468300782898]\n"
     ]
    }
   ],
   "source": [
    "reward_rate = 1e18\n",
    "num_tokens_per_pool = 5*[100_000_000*1e18]\n",
    "owner_fees_per_pool = [0] + 4*[500] # first pool owner takes no fee, other 4 owners take 5%\n",
    "print(\"APRs in % for each pool  :\" , compute_native_APR(reward_rate, num_tokens_per_pool, owner_fees_per_pool))\n",
    "print(\"APYs in % for same pools :\", compute_native_APY(reward_rate, num_tokens_per_pool, owner_fees_per_pool, 365))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9226d508-6b2a-470b-a7ca-10d0c6d3b170",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "APRs in % for each pool  : [73.07121951219511, 7.307121951219512, 7.307121951219512, 7.307121951219512, 7.307121951219512]\n",
      "APYs in % for same pools : [107.5042726922077, 7.579928500245625, 7.579928500245625, 7.579928500245625, 7.579928500245625]\n"
     ]
    }
   ],
   "source": [
    "reward_rate = 1e18\n",
    "num_tokens_per_pool = [1_000_000*1e18] + 4*[100_000_000*1e18] # first OperatorStaking pool has -99% less tokens staked than the other 4 pools\n",
    "owner_fees_per_pool = 5*[500] # first pool owner takes no fee, other 4 owners take 5%\n",
    "print(\"APRs in % for each pool  :\" , compute_native_APR(reward_rate, num_tokens_per_pool, owner_fees_per_pool))\n",
    "print(\"APYs in % for same pools :\", compute_native_APY(reward_rate, num_tokens_per_pool, owner_fees_per_pool, 365))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
